package com.xiao.framework.ds.autoconfigure;

import com.xiao.framework.ds.autoconfigure.bean.GroupNameDataSource;
import com.xiao.framework.ds.autoconfigure.context.DataSourceHolder;
import com.xiao.framework.ds.autoconfigure.factory.MultiDataSourceFactory;
import com.xiao.framework.ds.autoconfigure.strategy.MultiDataSourceDetermineStrategy;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import javax.sql.DataSource;
import lombok.Setter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.StringUtils;

/**
 * @author <a href="mailto:xiaoquanidea@163.com">hutiquan</a>
 * @since 2019-08-24 1:06
 */
public class MultiRoutingDataSource extends AbstractRoutingDataSource implements InitializingBean {
  private static final Logger log = LoggerFactory.getLogger(MultiRoutingDataSource.class);

  private static final String GROUP_KEY  = "_";

  private final MultiDataSourceFactory multiDataSourceFactory;

  private DataSource defaultPrimaryDataSource;

  private Map<String, DataSource> storeDataSourceMap = new LinkedHashMap<>();

  private Map<String, GroupNameDataSource> groupDataSourceMap = new ConcurrentHashMap<>();

  @Setter private String primary;

  @Setter private MultiDataSourceDetermineStrategy multiDataSourceDetermineStrategy;


  public MultiRoutingDataSource(MultiDataSourceFactory multiDataSourceFactory) {
    this.multiDataSourceFactory = multiDataSourceFactory;
  }

  @Override
  public DataSource determineDataSource() {
    String dataSourceKey = DataSourceHolder.get();
    if (!StringUtils.hasText(dataSourceKey)) {
      if (log.isDebugEnabled()) {
        log.debug("返回默认主数据源:{}",primary);
      }
      return defaultPrimaryDataSource;
    }
    if (groupDataSourceMap.containsKey(dataSourceKey)) {
      DataSource groupDataSource = groupDataSourceMap.get(dataSourceKey).demeterDataSource();
      if (log.isDebugEnabled()) {
        log.debug("从{}组中返回数据源",dataSourceKey);
      }
      return groupDataSource;
    }else if (storeDataSourceMap.containsKey(dataSourceKey)){
      DataSource dataSource = storeDataSourceMap.get(dataSourceKey);
      if (log.isDebugEnabled()) {
        log.debug("从数据源列表中返回数据源,数据源:{}",dataSourceKey);
      }
      return dataSource;
    }
    throw new IllegalStateException("没有获取到目标数据源,数据源:" + dataSourceKey);
  }

  @Override
  public void afterPropertiesSet() throws Exception {
    Map<String, DataSource> dataSourceMap = multiDataSourceFactory.loadDataSource();
    if (log.isInfoEnabled()) {
      log.info("初始数据源完成,共加载 {} 个,--------->:{}", dataSourceMap.size(),dataSourceMap.keySet().stream().collect(
          Collectors.joining(",","[","]"))
      );
    }
    storeDataSourceMap.putAll(dataSourceMap);
    this.defaultPrimaryDataSource = determinePrimaryDataSource();
    if (log.isInfoEnabled()) {
      log.info("当前主数据源为{}",primary);
    }
    dataSourceMap.forEach((key, dataSource) -> {
      if (key.contains(GROUP_KEY)) {
        String groupName = key.split(GROUP_KEY)[1];
        if (groupDataSourceMap.containsKey(groupName)) {
          groupDataSourceMap.get(groupName).addDataSource(dataSource);
        }else {
          GroupNameDataSource groupNameDataSource = new GroupNameDataSource(groupName, multiDataSourceDetermineStrategy);
          groupNameDataSource.addDataSource(dataSource);
          groupDataSourceMap.put(groupName, groupNameDataSource);
        }
      }
    });
  }

  public DataSource determinePrimaryDataSource() {
    if (storeDataSourceMap.containsKey(primary)) {
      return storeDataSourceMap.get(primary);
    }
    throw new IllegalStateException("没有配置primary数据源!");
  }

}
